package org.btik.light.tool;


import java.math.BigInteger;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;

/**
 * DH密钥协商工具，只用作密钥交换
 *
 * @author lustre
 * @since 2023/6/7 0:15
 */
public class DhUtil {
    /**
     * 非对称加密密钥算法
     */
    private static final String KEY_ALGORITHM = "DH";
    /**
     * 本地密钥算法，即对称加密密钥算法<br>
     * {@link com.sun.crypto.provider.DHKeyAgreement#engineGenerateSecret(String)}其中有相应的分支判断
     * 关于系统属性设置jdk.crypto.KeyAgreement.legacyKDF为true的原因也在其中
     */
    private static final String SELECT_ALGORITHM = "AES";

    /**
     * 密钥长度
     */
    private static final int KEY_SIZE = 512;


    /**
     * 初始化甲方密钥
     *
     * @return DhKeyPair 公私钥对
     * @throws NoSuchAlgorithmException 一般是支持DH算法的，正常情况下不会抛出，除非包缺失
     */
    public static DhKeyPair initKey() throws NoSuchAlgorithmException {
        //实例化密钥对生成器
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        //初始化密钥对生成器
        keyPairGenerator.initialize(KEY_SIZE);
        //生成密钥对
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        return new DhKeyPair(keyPair);
    }

    /**
     * 初始化乙方密钥 ， 生成私钥并使用对端的 P G值和私钥计算公钥
     *
     * @param encodeDhPublicKey 对端编码过的公钥 含 G P 值和对端的 公钥
     */
    public static DhKeyPair initKey(byte[] encodeDhPublicKey) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeySpecException {
        // 将编码过的公钥转回DHPublicKey类型，获取其中的P G值和长度
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(encodeDhPublicKey);
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        PublicKey pubKey = keyFactory.generatePublic(x509KeySpec);
        return initKey(((DHPublicKey) pubKey).getParams());
    }

    /**
     * 初始化乙方密钥 ， 生成私钥并使用对端的 P G值和私钥计算公钥
     * B = G ^ privateKeyB mode P
     *
     * @param dhParameterSpec 使用 P G值 计算公钥
     */
    public static DhKeyPair initKey(DHParameterSpec dhParameterSpec) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
        //实例化密钥对生成器
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KEY_ALGORITHM);
        //初始化密钥对生成器
        keyPairGenerator.initialize(KEY_SIZE);
        keyPairGenerator.initialize(dhParameterSpec);
        //产生密钥对
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        return new DhKeyPair(keyPair);
    }

    /**
     * 协商出密钥<br>
     * 该方式适用于对端也是java的情况。
     * 这里的对端公钥和本地私钥都使用{@link Key#getEncoded()}来编码key.<br>
     *
     * @param publicKey  对端公钥
     * @param privateKey 本端私钥
     * @return 协商出的密钥   K = 对端公钥的X值 ^ 本端私钥的X值 mode P 当然这里公钥含了 P G 信息
     */
    public static byte[] keyAgree(byte[] publicKey, byte[] privateKey) throws Exception {
        //实例化密钥工厂
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        //初始化公钥
        //密钥材料转换
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(publicKey);
        //产生公钥
        PublicKey pubKey = keyFactory.generatePublic(x509KeySpec);
        //初始化私钥
        //密钥材料转换
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(privateKey);
        //产生私钥
        PrivateKey priKey = keyFactory.generatePrivate(pkcs8KeySpec);
        //实例化
        KeyAgreement keyAgree = KeyAgreement.getInstance(keyFactory.getAlgorithm());
        //初始化
        keyAgree.init(priKey);
        keyAgree.doPhase(pubKey, true);
        //生成本地密钥
        SecretKey secretKey = keyAgree.generateSecret(SELECT_ALGORITHM);
        return secretKey.getEncoded();
    }

    /**
     * 协商出密钥<br>
     * 该方式适用于对端也是java的情况。对端公钥也使用{@link Key#getEncoded()}来编码key.<br>
     *
     * @param publicKey    对端公钥
     * @param dhPrivateKey 本端私钥
     * @return 协商出的密钥   K = 对端公钥的X值 ^ 本端私钥的X值 mode P 当然这里公钥含了 P G 信息
     */
    public static byte[] keyAgree(byte[] publicKey, DHPrivateKey dhPrivateKey) throws Exception {
        //实例化密钥工厂
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        // 将公钥编码转公钥
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(publicKey);
        PublicKey pubKey = keyFactory.generatePublic(x509KeySpec);

        //完成协商
        KeyAgreement keyAgree = KeyAgreement.getInstance(keyFactory.getAlgorithm());
        keyAgree.init(dhPrivateKey);
        keyAgree.doPhase(pubKey, true);
        SecretKey secretKey = keyAgree.generateSecret(SELECT_ALGORITHM);
        return secretKey.getEncoded();
    }

    /**
     * 协商出密钥<br>
     * 适用于对端只返回对端公钥数值，不能返回适用于java的公钥编码时使用。
     *
     * @param y            对端公钥
     * @param dhPrivateKey 本端私钥
     * @return 协商出的密钥   K = 对端公钥的X值 ^ 本端私钥的X值 mode P 当然这里公钥含了 P G 信息
     */
    public static byte[] keyAgree(BigInteger y, BigInteger p, BigInteger g, DHPrivateKey dhPrivateKey) throws Exception {
        KeyFactory keyFactory = KeyFactory.getInstance(KEY_ALGORITHM);
        //实例化
        KeyAgreement keyAgree = KeyAgreement.getInstance(keyFactory.getAlgorithm());
        //初始化
        keyAgree.init(dhPrivateKey);
        keyAgree.doPhase(new DHPublicKeyImpl(y, p, g), true);
        //生成本地密钥
        SecretKey secretKey = keyAgree.generateSecret(SELECT_ALGORITHM);
        return secretKey.getEncoded();
    }

    public static class DhKeyPair {
        DHPublicKey publicKey;

        DHPrivateKey privateKey;

        public DHPublicKey getPublicKey() {
            return publicKey;
        }

        public DHPrivateKey getPrivateKey() {
            return privateKey;
        }

        public DhKeyPair(KeyPair keyPair) {
            this.publicKey = (DHPublicKey) keyPair.getPublic();
            this.privateKey = (DHPrivateKey) keyPair.getPrivate();
        }
    }

    /**
     * 对端的公钥编码不一定是java的编码方式，可能不含长度和P G值，这里需要使用对端的公钥和 本地保存的P G来构建公钥
     */
    static class DHPublicKeyImpl implements DHPublicKey {
        private final BigInteger y;

        private final DHParameterSpec params;

        @Override
        public BigInteger getY() {
            return y;
        }

        @Override
        public String getAlgorithm() {
            return KEY_ALGORITHM;
        }

        @Override
        public String getFormat() {
            return "X.509";
        }

        @Override
        public byte[] getEncoded() {
            return new byte[0];
        }

        @Override
        public DHParameterSpec getParams() {
            return params;
        }


        public DHPublicKeyImpl(BigInteger y, BigInteger p, BigInteger g) {
            this.y = y;
            this.params = new DHParameterSpec(p, g);
        }
    }

}